package com.ng8866.gateway.filter;


import com.ng8866.gateway.balancer.LeastConnectionLoadBalancer;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;

import java.net.URI;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

@Slf4j
@Order(100)
@Component
@ConditionalOnProperty(prefix = "loadBalancer", name = "enableBalancer", havingValue = "true")
public class LeastConnectionGlobalFilter implements GlobalFilter {
    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        return chain.filter(exchange).doFinally(type -> {
            log.debug("===> 请求处理完成");
            // 获取路由信息
            Route route = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);
            if (route != null) {
                String routeAuthority = route.getUri().getAuthority();
                log.debug("===> 路由 authority：" + routeAuthority);
                processRouteAuthority(exchange, routeAuthority);
            } else {
                log.warn("===> 未能获取到路由信息");
            }
            log.debug("===> 请求处理完成{}", LeastConnectionLoadBalancer.loadMap);
        });
    }

    private void processRouteAuthority(ServerWebExchange exchange, String routeAuthority) {
        ConcurrentHashMap<String, AtomicInteger> routeMap = LeastConnectionLoadBalancer.loadMap.get(routeAuthority);
        log.debug("===> 路由连接数：" + routeMap);
        if (routeMap != null && !routeMap.isEmpty()) {
            log.debug("===> 路由连接数大于 0");
            URI uri = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR);
            log.debug("===> 请求转发的服务器信息：" + uri);
            if (uri != null) {
                String requestAuthority = uri.getAuthority();
                log.debug("===> 请求转发的服务器 authority：" + requestAuthority);
                decrementConnectionCount(routeMap, requestAuthority);
                LeastConnectionLoadBalancer.setCountMap(requestAuthority);
                log.debug("===> 总数" + LeastConnectionLoadBalancer.countMap);
            } else {
                log.warn("===> 未能获取到请求转发的服务器信息");
            }
        } else {
            log.debug("===> 路由连接数不足或不存在对应路由的连接数信息");
        }
    }

    private void decrementConnectionCount(ConcurrentHashMap<String, AtomicInteger> routeMap, String requestAuthority) {
        if (routeMap == null) {
            log.info("Route map is null.");
            return;
        }
        routeMap.forEach((key, count) -> {
            log.debug("===> 路由 key：" + key);
            log.debug("===> 路由连接数（减前）：" + count);
            if (key.contains(requestAuthority) && count != null) {
                int oldValue = count.getAndDecrement();
                log.info("===> 路由 key：" + key + "，连接数（减后）：" + count.get());
                if (oldValue < 0) {
                    count.set(0);
                }
            }
        });
    }


}
